package cn.jbolt.ai.core.fileSplit;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.core.AIPrompts;
import cn.jbolt.ai.klb.DocDataType;
import cn.jbolt.ai.klb.FileType;
import cn.jbolt.ai.klb.entity.AiDoc;
import cn.jbolt.ai.klb.entity.AiFile;
import cn.jbolt.ai.klb.service.AiDocService;
import cn.jbolt.util.*;
import com.jboltai.capability.model.JBoltAI;
import com.jboltai.capability.text.JBoltText;
import com.jboltai.event.ai.aichat.AIChatEvent;
import com.jboltai.event.text.TextSplitEvent;
import com.jboltai.resource.ai.AIModel;
import com.jboltai.util.text.FileInfo;
import com.jboltai.util.text.TextUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.InputStream;
import java.io.RandomAccessFile;
import java.nio.charset.Charset;
import java.util.List;
import java.util.function.Consumer;

/**
 *
 */
@Service
public class FileSplitServiceImpl implements FileSplitService {

    private static final Logger logger = LoggerFactory.getLogger(FileSplitServiceImpl.class);
    @Autowired
    private QiniuUtil qiniuUtil;
    @Autowired
    private UploadLocalUtil uploadLocalUtil;
    @Autowired
    private FilePathGeneratorUtil filePathGeneratorUtil;

    @Autowired
    private AiDocService docService;


    @Override
    public Result<String> split(AiFile file) {

        if (file.getDocType().equals(FileType.NORMAL.getCode())) {
            return splitNormalFile(file);
        }


        return splitQAFile(file);


    }

    /**
     * @param aiFile
     * @return
     */
    private Result splitQAFile(AiFile aiFile) {
        Result<File> result = splitByChunk(aiFile.getFilePath(), 2000, 0, aiFile.getExtractResources(),
                aiFile.getResourcesPosition(), aiFile.getPath());
        if (result.isError()) {
            return Result.error(result.getMsg());
        }
        try {
            FileSplitCache.addFileSplitCounter(aiFile.getId());
            StringBuilder aiProcessingDoc = new StringBuilder();
            //开始解析
            result = processBlocks(result.getData(), (block) -> {
                //拿到这个块，要继续让AI提取QA的
                AIChatEvent event = aiSplitText(aiProcessingDoc, block, aiFile);
                List<String> blocks = event.getAttr("blocks");
                extractQAFromBlocks(aiFile, blocks);

            });
            if (result.isError()) {
                return result;
            }
            if (aiProcessingDoc.length() > 0) {
                //还有剩余的文本
                extractQAFromBlock(aiFile, aiProcessingDoc.toString());
            }
            return result;
        } catch (Exception e) {
            logger.error("AI切分文件：{}, 失败", aiFile.getFilePath());
            return Result.error(e.getMessage());
        }
    }


    /**
     * AI
     *
     * @param doc
     */
    private AIChatEvent aiSplitText(StringBuilder processing, String doc, AiFile aiFile) {
        /*
          doc是 通过jbolttext根据token数分割的文本段，这个文本段接下来要交给ai 根据语义进行分割。
          分割的最后一部分，可能不是完整的，就留在processing中，和下一个 token文本段合并再交给ai处理。
         */
        processing.replace(0, processing.length(), mergeStrings(processing.toString(), doc));
        AIChatEvent event = JBoltAI.chat().prompt(AIPrompts.SPLIT_TEMPLATE, processing);
        if (aiFile.getAiModel() != null) {
            event.setModel(aiFile.getAiModel());
        }
        if (StrUtil.isNotEmpty(aiFile.getDefaultModel())) {
            event.setModelName(aiFile.getDefaultModel()); //设置默认模型
        }
        event.setAttr("oriText", processing);
        event.setStream(false).setTemperature(0.1)
                .setMaxTokens(2000)
                .setHandler(new AISplitHandler(aiFile))
                .publish().await();

        if (event.isFail()) {
            throw new RuntimeException(event.getError().getMsg(), event.getError().getException());
        }

        return event;
    }

    /**
     * ai根据语义将文本分成了多个文本块，接下来将文本块交给大模型提取题目，但文本块可能太小，要凑500个token大小再去提取
     *
     * @param file
     * @param splitBlocks
     */
    private void extractQAFromBlocks(AiFile file, List<String> splitBlocks) {
        double tokenSum = 0;
        double itemToken = 0;
        StringBuilder tempStr = new StringBuilder();

        for (int i = 0; i < splitBlocks.size(); i++) {
            itemToken = TextUtil.calToken(splitBlocks.get(i));
            if (tokenSum + itemToken > 1800 && tokenSum > 0) {
                //满2000了，提取一次
                extractQAFromBlock(file, tempStr.toString());
                tempStr.setLength(0);
                tokenSum = 0;
            }
            tokenSum += itemToken;
            tempStr.append(splitBlocks.get(i));

        }
        if (tokenSum > 0) {
            extractQAFromBlock(file, tempStr.toString());
        }
    }

    private void extractQAFromBlock(AiFile file, String block) {
        AIChatEvent event = JBoltAI.chat().prompt(AIPrompts.EXTRACT_QA_TEMPLATE, block)
                .setStream(false).setHandler(new AIExtractQAHandler(file))
                .setMaxTokens(2000);
        if (file.getAiModel() != null) {
            event.setModel(file.getAiModel());
        }
        if (StrUtil.isNotEmpty(file.getDefaultModel())) {
            event.setModelName(file.getDefaultModel()); //设置默认模型
        }
        event.publish();
        event.await();
        if (event.isFail()) {
            throw new RuntimeException(event.getError().getMsg(), event.getError().getException());
        }

        List<AiDoc> docs = event.getAttr("docs");
        docs.forEach(doc -> {
            FileSplitCache.addTotalCount(file.getId());
            docService.add(doc);
        });
    }


    /**
     * @param a
     * @param b
     * @return
     */
    public static String mergeStrings(String a, String b) {
        if (StrUtil.isBlank(a)) {
            return b;
        }
        if (StrUtil.isBlank(b)) {
            return a;
        }
        int maxOverlap = findMaxOverlap(a, b);
        return a + b.substring(maxOverlap);
    }

    private static int findMaxOverlap(String a, String b) {
        int low = 0;
        int high = Math.min(a.length(), b.length());
        int maxOverlap = 0;

        while (low <= high) {
            int mid = (low + high) / 2;
            if (a.endsWith(b.substring(0, mid))) {
                maxOverlap = mid;
                low = mid + 1;
            } else {
                high = mid - 1;
            }
        }
        return maxOverlap;
    }


    private Result splitNormalFile(AiFile aiFile) {
        if (aiFile.getChunkSize() == null || aiFile.getChunkSize() == 0) {
            return Result.error("请输入分块数量");
        }
        if (aiFile.getChunkOverlap() == null || aiFile.getChunkOverlap() == 0) {
            return Result.error("请输入重叠数量");
        }
        Result<File> result = splitByChunk(aiFile.getFilePath(), aiFile.getChunkSize(), aiFile.getChunkOverlap(),
                aiFile.getExtractResources(), aiFile.getResourcesPosition(), aiFile.getPath());
        if (result.isError()) {
            return Result.error(result.getMsg());
        }

        FileSplitCache.addFileSplitCounter(aiFile.getId());

        //开始解析
        result = processBlocks(result.getData(), (block) -> {
            FileSplitCache.addTotalCount(aiFile.getId());
            AiDoc doc = new AiDoc();
            doc.setTitle(block);
            doc.setKnowledgeBaseId(aiFile.getKnowledgeBaseId());
            doc.setFileId(aiFile.getId());
            doc.setDataType(DocDataType.FILE_PIECE.getCode());
            doc.setDocType(aiFile.getDocType());
            docService.add(doc);
        });

        return result;

    }

    /**
     * @param file
     * @param textConsumer
     */
    private Result processBlocks(File file, Consumer<String> textConsumer) {
        Result result = Result.success();
        StringBuilder cache = new StringBuilder();
        try (RandomAccessFile accessFile = new RandomAccessFile(file, "r")) {
            FileUtil.readLines(accessFile, Charset.forName("utf-8"), (line) -> {
                if (result.isError()) return;
                try {
                    if (FilesUtil.splitSymbol.equals(line)) {
                        //上一个块结束了
                        textConsumer.accept(cache.toString());
                        cache.delete(0, cache.length());
                    } else {
                        cache.append(line);
                    }
                } catch (Exception e) {
                    result.setState(Result.ERROR_STATE);
                }

            });
            if (cache.length() > 0) {
                textConsumer.accept(cache.toString());
                cache.delete(0, cache.length());
            }
            return result;
        } catch (Exception e) {
            logger.error("文件处理失败，失败原因：" + e.getMessage());
            return Result.error("文件处理失败");
        } finally {
            FileUtil.del(file);
        }

    }


    /**
     * @param url
     * @param chunkSize
     * @param extractResources  是否提取文档中的资源,图片,音频之类的,如果是true,就进onfile
     * @param resourcesPosition 资源存储路径,本地或者qiniu
     * @param filePath          存储的目录地址
     */
    public Result<File> splitByChunk(String url, Integer chunkSize, Integer chunkOverlap, Boolean extractResources,
                                     String resourcesPosition, String filePath) {
        if (StrUtil.isBlankIfStr(url)) {
            return Result.error("请输入要提取的文档地址");
        }
        //构建一个临时文件存储分割的内容
        Result<File> result = FilesUtil.createFileInFolder(RandomUtil.randomString(8) + ".txt");
        if (result.isError()) {
            return result;
        }
        File tmpFile = result.getData();
        try {
            TextSplitEvent event = configureTextSplitEvent(url, chunkSize, chunkOverlap);
            event.onSuccess((e, res) -> {
                String content = res.getChuncks().get(0) + "\n" + FilesUtil.splitSymbol + "\n";
                logger.info("文档分割中，分割到内容:{}", content);
                FileUtil.appendUtf8String(content, tmpFile);
                logger.info("文档分割中内容已写入临时文件： {}", tmpFile.getAbsolutePath());
            });
            event.onFail((e, error) -> {
                logger.error("文档分割失败，文档地址为:{},失败原因:{}", url, error.getFullMsg());
            });

            if (extractResources!=null&&extractResources) {
                event.onFile((e, fileInfo) -> {

                     if(resourcesPosition.equals(AiFile.RESOURCES_QINIU)){
                         return uploadQiniu(fileInfo, filePath);
                     }else {
                         return uploadLocal(fileInfo, filePath);
                     }
                });
            }

            event.publish().await();

            if (event.isFail()) {
                logger.error("文档分割失败，文档地址为:{},失败原因:{}", event.getUrl(), event.getError().getMsg());
                tmpFile.deleteOnExit();
                return Result.error(event.getError().getMsg());
            }


            return Result.success(tmpFile);
        } catch (Exception e) {
            logger.error("文档分割失败，文档地址为:{},失败原因:{}", url, e.getMessage());
            return Result.error(e.getMessage());
        }
    }

    /**
     * onfile中上传资源资源到本地
     * @param fileInfo
     * @param filePath
     * @return
     */
    private String uploadLocal(FileInfo fileInfo, String filePath) {
        try {
            String fileName = fileInfo.getFileName();
            if(StrUtil.isNotEmpty(filePath)){
                fileName = filePathGeneratorUtil.generateFilePath(filePath, fileName);
            }
            //构建一个filePath的名字解析
            fileInfo.getInputStream();
            Result<String> fileRes = uploadLocalUtil.uploadInputStreamToLocal(fileInfo.getInputStream(), fileName);
            if (fileRes.isSuccess()) {
                String path = fileRes.getData();
                logger.info("文档提取的文件上传成功, 文件地址: {}", path);
                return "![" + fileName + "](" + path + ")";
            } else {
                return "![" + fileName + "]('')";
            }
        } catch (Exception error) {
            String fileName = fileInfo.getFileName();
            logger.error("文档提取的文件上传失败,文件名: {}, 失败原因: {}", fileName, error.getMessage());
            return "![" + fileName + "]('')";
        }
    }

    /**
     * onfile中上传资源到七牛
     * @param fileInfo
     * @param filePath
     * @return
     */
    private String uploadQiniu(FileInfo fileInfo, String filePath) {
        try {
            String fileName = fileInfo.getFileName();
            if(StrUtil.isNotEmpty(filePath)){
                filePath = filePathGeneratorUtil.generateFilePath(filePath, fileName);
            }
            Result<String> fileRes = qiniuUtil.uploadFile(fileInfo.getInputStream(), fileName, filePath);
            if (fileRes.isSuccess()) {
                String path = fileRes.getData();
                logger.info("文档提取的文件上传成功, 文件地址: {}", path);
                return "![" + fileName + "](" + path + ")";
            } else {
                return "![" + fileName + "]('')";
            }
        } catch (Exception error) {
            String fileName = fileInfo.getFileName();
            logger.error("文档提取的文件上传失败,文件名: {}, 失败原因: {}", fileName, error.getMessage());
            return "![" + fileName + "]('')";
        }
    }


    /**
     *
     */
    private TextSplitEvent configureTextSplitEvent(String url, Integer chunkSize, Integer chunkOverlap) {
        TextSplitEvent textSplitEvent = JBoltText.splitUrl(url).setStream(true);
        if (chunkSize != null && chunkSize > 0) {
            textSplitEvent.setChunkSize(chunkSize);
        }
        if (chunkOverlap != null && chunkOverlap >= 0) {
            textSplitEvent.setChunkOverlap(chunkOverlap);
        }
        return textSplitEvent;
    }


}
